#include <windows.h>
#include "config.h"
#include <vector>
#include <iostream>
#include <sstream>
#include <fstream>
#include <string>

#include "MutualInformation.h"
#include "Scale.h"

__int64 FeatureNo;
__int64 SampleNo;

// mRMR schema: MIQ (mutual information quotient) or MID (mutual information difference).
// MIQ seems to perform better nearly always.
#define MIQ 1
#define MID 2
#define SCHEMA MIQ

// Number of threads to use
// Values other than 8 has not been tested yet, so i don't recoommend using
#define CORES 8

using namespace std;

void ImportTransposedFile(short * classes, short * features, __int64 numSamples, __int64 numFeatures, ifstream &fClasses) {
	// Second line of file: classes for all samples
	for (int sample = 0; sample < numSamples; ++sample) {
		fClasses >> classes[sample];
	}

	__int64 sf = numSamples*numFeatures;
	// Load values int feature matrix
	for (__int64 feature = 0; feature < sf; ++feature) {
		fClasses >> features[feature];
	}
}


// this assumes little-endian processor bit order [http://en.wikipedia.org/wiki/Endianness]
struct bits
{
     unsigned int bit8 : 1;
	 unsigned int bit7 : 1;
	 unsigned int bit6 : 1;
	 unsigned int bit5 : 1;
	 unsigned int bit4 : 1;
	 unsigned int bit3 : 1;
	 unsigned int bit2 : 1;
	 unsigned int bit1 : 1;
};


void ImportBinary(short * classes, short * features, __int64 numSamples, __int64 numFeatures, ifstream &fSamples, ifstream &fClasses) {
	// cout << "importing " << numSamples << " samples, features: " << numFeatures << endl;

	// Second line of file: classes for all samples
	for (int sample = 0; sample < numSamples; ++sample) {
		fClasses >> classes[sample];
	}

	__int64 sf = numSamples*numFeatures;
	__int64 numChunks = ceil(sf / 8.0);

	// Load values int feature matrix
	for (__int64 chunk = 0; chunk < numChunks; ++chunk) {
		unsigned char character = fSamples.get();
		struct bits * num = (bits*) &character;
		
		features[chunk*8 + 0] = num->bit1;
		features[chunk*8 + 1] = num->bit2;
		features[chunk*8 + 2] = num->bit3;
		features[chunk*8 + 3] = num->bit4;
		features[chunk*8 + 4] = num->bit5;
		features[chunk*8 + 5] = num->bit6;
		features[chunk*8 + 6] = num->bit7;
		features[chunk*8 + 7] = num->bit8;
	}

}

/* features vector sholud be at least 8 bytes longer than number of features,
 * for additional bits might be loaded that fill a byte at the end*/
void ImportSamplesFromFile(string fileName, __int64 numFeatures, short* features) {
	__int64 numChunks = ceil(numFeatures / 8.0);

	ifstream file (fileName, ios::in|ios::binary);

	for (__int64 chunk = 0; chunk < numChunks; ++chunk) {
		unsigned char character = file.get();
		struct bits * num = (bits*) &character;
		
		features[chunk*8 + 0] = num->bit1;
		features[chunk*8 + 1] = num->bit2;
		features[chunk*8 + 2] = num->bit3;
		features[chunk*8 + 3] = num->bit4;
		features[chunk*8 + 4] = num->bit5;
		features[chunk*8 + 5] = num->bit6;
		features[chunk*8 + 6] = num->bit7;
		features[chunk*8 + 7] = num->bit8;
	}

}


int getMaxIndex(float * arrayData, int arrayLength, int* selectable) {
	float max = arrayData[0];
	int maxIndex = 0;

	for (int i = 1; i < arrayLength; ++i) {
		if ((arrayData[i] > max) && (selectable[i] > 0)) { maxIndex = i; max = arrayData[i]; }
	}

	return maxIndex;
}

struct thrOptions {
	__int64 startColumn, endColumn, lastAddedIndex;
	float * redundancy, * score, * miOnTargetClass, number;
	short * features;
	int * selectable;
	int ignore;
	int threadIdx;
};

void calculateScores(LPVOID o) {
	thrOptions * opt = (thrOptions *) o;
	if (opt->ignore) return; // nothing to do on this core

	for (__int64 k = opt->startColumn; k < opt->endColumn; ++k) {
		int selectable = opt->selectable[k];

		if (selectable) {
			short * lastFeatureAddress = opt->features + SampleNo * opt->lastAddedIndex;
			short * checkedAddress = opt->features + SampleNo * k;

			double mi = calculateMutualInformation(lastFeatureAddress, checkedAddress, SampleNo, opt->threadIdx);
			opt->redundancy[k] += mi;
			
			// try redundancy reduced by 2.5-2?

#if SCHEMA == MIQ
				double score = opt->miOnTargetClass[k] / ((opt->redundancy[k] / opt->number) + 0.01); // MIQ, better results in iris feat selection
#else
				double score = ( opt->miOnTargetClass[k] - opt->redundancy[k] / opt->number ); // MID scheme
#endif
			opt->score[k] = score;
		} else {
			opt->score[k] = 0;
		}

	}
}


void normaliseVector(short *vector, int vectorLength)
{
	short minVal = vector[0];
	short maxVal = vector[0];

	for (int i = 1; i < vectorLength; ++i) {
		if (vector[i] < minVal)	minVal = vector[i];
		if (vector[i] > maxVal)	maxVal = vector[i];
	}
      
	for (int i = 0; i < vectorLength; ++i) {
		vector[i] -= minVal;
	}
	maxVal -= minVal;

	// Optimize reducing max state number
	int * mapping = (int*) calloc(maxVal + 1, sizeof(int));

	for (int i = 0; i < vectorLength; ++i) {
		mapping[vector[i]] = 1;
	}

	int last = 0;
	for (int i = 0; i < maxVal + 1; ++i) {
		if (mapping[i] == 1) mapping[i] = last++;
	}

	// Apply mapping
	for (int i = 0; i < vectorLength; ++i) {
		vector[i] = mapping[vector[i]];	
	}

	free(mapping);
}


void mRMR_CPU(short * features, short * classes, int featuresToSelectNo) {
	// Prepare memory for calculation results
	MemInit();

	// Mutual information between target class and i-th feature
	float * miOnTargetClass = new float[FeatureNo];
	int * selectable = new int[FeatureNo];
	for (int i = 0; i < FeatureNo; ++i) selectable[i] = 1;

	// Temporary array to store scores for different features
	float * redundancy = new float[FeatureNo];
	for (int i = 0; i < FeatureNo; ++i) redundancy[i] = 0;

	float * score = new float[FeatureNo];

	// mRMR result list
	int * resultSet_FeatureIndex = new int[featuresToSelectNo];
	float * resultSet_FeatureScore = new float[featuresToSelectNo];
	
	//cout << "Normalizing classes" << endl;
	normaliseVector(classes, SampleNo);
	//for (int i = 0; i < SampleNo; ++i) cout << i << ":" << classes[i] << " ";

	//cout << "Normalizing features" << endl;
	// Prepare input data - normalize vectors
	//for (int i = 0; i < FeatureNo; ++i) {
	//	normaliseVector(features + i*SampleNo, SampleNo);
	//} no need for binary data

	//cout << endl << "Calculating MI with target class" << endl;
	// Count MI between target class and each feature
	for (__int64 i = 0; i < FeatureNo; ++i) {
		miOnTargetClass[i] = calculateMutualInformation(classes, features + i * SampleNo, SampleNo, 0);
		
		//cout << "MI #" << i << " is " << miOnTargetClass[i] << endl;
	}

	// Select the features with minimum redundancy
	
	// First feature is the one with maximum MI with the terget class
	//cout << "Finding first feature with max MI" << endl;
	int firstIndex = getMaxIndex(miOnTargetClass, FeatureNo, selectable);
	resultSet_FeatureIndex[0] = firstIndex;
	resultSet_FeatureScore[0] = miOnTargetClass[ firstIndex ]; 
	selectable[firstIndex] = 0;

	
	// The index of last added feature to recalculate redundancy (add MI with this feature to sum)
	__int64 lastAddedIndex = firstIndex;

	// Output result immediately
	cout << (lastAddedIndex + 1) << " " << miOnTargetClass[ lastAddedIndex ] << endl;

	// Prepare for threading
	int ftPerThread = FeatureNo / CORES + 1;
	HANDLE handles[CORES];
	thrOptions opt[CORES];

	for (int s=0; s<CORES; ++s) {
		opt[s].threadIdx = s;
		opt[s].features = features;
		opt[s].lastAddedIndex = lastAddedIndex;
		opt[s].miOnTargetClass = miOnTargetClass;
		opt[s].redundancy = redundancy;
		opt[s].score = score;
		opt[s].selectable = selectable;

		opt[s].startColumn = ftPerThread * s;
		opt[s].endColumn = ftPerThread * s + ftPerThread;
		opt[s].ignore = 0;

		if (opt[s].endColumn > FeatureNo) opt[s].endColumn = FeatureNo;
		if (opt[s].startColumn > FeatureNo) opt[s].ignore = 1;
		
		opt[s].number = 0;
	}
	
	// Select featuresToSelectNo features
	for (int i = 1; i < featuresToSelectNo; ++i) {
		float floatI = (float) i;
		// Calculate score for each 'selectable' feature

		// use as many threads as defined in const value
		for (int t = 0; t < CORES; ++t) {
			opt[t].number = floatI; // number of already-selected features
			opt[t].lastAddedIndex = lastAddedIndex; // important!

			handles[t] = CreateThread(NULL, //Choose default security
							0, //Default stack size
							(LPTHREAD_START_ROUTINE)& calculateScores, //Routine to execute
							(LPVOID) (opt + t), // Thread parameter
							0, // Immediately run the thread
							NULL); // Thread Id
		}

		// --- Time for threads to execute --- //

		// Join threads
		for (int t = 0; t < CORES; ++t) WaitForSingleObject(handles[t], INFINITE); // no scheme
		
		// Select feature with maximum score
		lastAddedIndex = getMaxIndex(score, FeatureNo, selectable);
		selectable[lastAddedIndex] = 0;

		// Store results in result array
		resultSet_FeatureIndex[i] = lastAddedIndex;
		resultSet_FeatureScore[i] = score[lastAddedIndex];

		// Output result immediately
		cout << (lastAddedIndex + 1) << " " << score[lastAddedIndex] << endl;
		// Indexes start with 1, not 0, for consistency with original matlab implementation's output
	}

	// Clean up mess
	delete miOnTargetClass;
	delete selectable;
	delete redundancy;
	delete score;
	delete resultSet_FeatureIndex;
	delete resultSet_FeatureScore;

}


typedef vector<string> CommandLineStringArgs;

void DisplayHelp() {
	cout << "Wrong number of parameters. See documentation for correct usage." << endl;
}

int main(int argc, char * argv[])
{
	if (argc == 1) {
		// wrong number of arguments assume user didn't read documentation
		DisplayHelp();
		return -1;
	}

	// args to string vector
	CommandLineStringArgs cmdlineStringArgs(&argv[0], &argv[0 + argc]);
	
	string classesFileName = cmdlineStringArgs[1];
	ifstream fSamples;
	if (argc > 2) {
		string samplesFileName = cmdlineStringArgs[2];
		fSamples.open(samplesFileName, ios::in|ios::binary);
	}
	// create input streams for data
	ifstream fClasses (classesFileName, ios::in);

	// read number of samples and features to correctly interprete data
	fClasses >> SampleNo >> FeatureNo;

	// allocate some memory
	short * Classes = new short [SampleNo];
	__int64 memToAlloc = SampleNo * (FeatureNo + 8) * sizeof(short);
	void * space = malloc(memToAlloc);
	
	if (space == 0) {
		// memory could not be allocated (need more RAM?)
		cout << "Space address: " << space << " Exiting." << endl;
		cout << "Memory could not be allocated. Do you have enough memory available?" << endl;
		return 0;
	}
	short * Features = (short *) space;
		
	if (argc == 2) {
		// 1 argument: assume data is in text file
		ImportTransposedFile(Classes, Features, SampleNo, FeatureNo, fClasses);
	} else if (argc == 3) {
		// 2 arguments: assume data in binary format
		ImportBinary(Classes, Features, SampleNo, FeatureNo, fSamples, fClasses);
	} else {
		// wrong number of arguments assume user didn't read documentation
		DisplayHelp();
		return -1;
	}

	int ToSelect = FeatureNo;
	
	// Run selection using CPU (GPU version is no longer maintained)
	mRMR_CPU(Features, Classes, ToSelect);
	
	return 0;
}
